import sqlite3
import os
import sys
import glob

def find_database_files():
    """查找可能的数据库文件"""
    backend_dir = os.path.dirname(os.path.dirname(__file__))
    possible_paths = [
        os.path.join(backend_dir, 'app', 'database', 'app.db'),
        os.path.join(backend_dir, 'app.db'),
        os.path.join(backend_dir, 'database.db'),
        os.path.join(backend_dir, 'instance', 'app.db')
    ]
    
    # 添加所有 .db 文件
    possible_paths.extend(glob.glob(os.path.join(backend_dir, '**', '*.db'), recursive=True))
    
    # 过滤出存在的文件
    existing_paths = [path for path in possible_paths if os.path.exists(path)]
    return existing_paths

def fix_ai_models_in_db(db_path):
    """在指定的数据库中修复 ai_models 表"""
    print(f"尝试修复数据库: {db_path}")
    
    # 连接数据库
    conn = sqlite3.connect(db_path)
    conn.row_factory = sqlite3.Row  # 使结果可以通过列名访问
    cursor = conn.cursor()
    
    # 检查表是否存在
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='ai_models'")
    if not cursor.fetchone():
        print(f"数据库 {db_path} 中没有 ai_models 表，跳过。")
        conn.close()
        return False
    
    print(f"在数据库 {db_path} 中找到 ai_models 表。")
    
    # 检查字段是否已存在
    cursor.execute("PRAGMA table_info(ai_models)")
    columns = [column[1] for column in cursor.fetchall()]
    
    # 添加 stream_enabled 字段（如果不存在）
    if 'stream_enabled' not in columns:
        print(f"在数据库 {db_path} 中添加 stream_enabled 字段...")
        try:
            cursor.execute("ALTER TABLE ai_models ADD COLUMN stream_enabled INTEGER NOT NULL DEFAULT 1")
            conn.commit()
            print(f"在数据库 {db_path} 中成功添加 stream_enabled 字段！")
        except sqlite3.Error as e:
            print(f"添加字段失败: {e}")
            
            # 尝试创建新表并迁移数据
            print("尝试通过创建新表并迁移数据来添加字段...")
            try:
                # 获取表结构
                cursor.execute("PRAGMA table_info(ai_models)")
                columns_info = cursor.fetchall()
                
                # 创建新表
                cursor.execute("""
                CREATE TABLE ai_models_new (
                    id TEXT PRIMARY KEY,
                    name TEXT NOT NULL,
                    provider TEXT NOT NULL,
                    model_id TEXT NOT NULL,
                    api_type TEXT NOT NULL,
                    description TEXT,
                    max_tokens INTEGER NOT NULL DEFAULT 4096,
                    temperature_default REAL NOT NULL DEFAULT 0.7,
                    api_key TEXT NOT NULL,
                    api_endpoint TEXT,
                    api_version TEXT,
                    stream_enabled INTEGER NOT NULL DEFAULT 1,
                    is_active INTEGER NOT NULL DEFAULT 1,
                    created_at TIMESTAMP NOT NULL,
                    updated_at TIMESTAMP NOT NULL
                )
                """)
                
                # 获取旧表中的所有列名（除了新添加的）
                old_columns = ", ".join([col[1] for col in columns_info])
                
                # 迁移数据
                cursor.execute(f"""
                INSERT INTO ai_models_new (
                    {old_columns}, stream_enabled
                )
                SELECT {old_columns}, 1 FROM ai_models
                """)
                
                # 删除旧表
                cursor.execute("DROP TABLE ai_models")
                
                # 重命名新表
                cursor.execute("ALTER TABLE ai_models_new RENAME TO ai_models")
                
                conn.commit()
                print("成功通过表重建添加了 stream_enabled 字段！")
            except sqlite3.Error as e2:
                print(f"表重建失败: {e2}")
                conn.rollback()
                conn.close()
                return False
    else:
        print(f"数据库 {db_path} 中已存在 stream_enabled 字段。")
    
    # 获取所有 AI 模型记录
    cursor.execute("SELECT * FROM ai_models")
    models = cursor.fetchall()
    
    if not models:
        print(f"数据库 {db_path} 中没有找到 AI 模型记录。")
    else:
        print(f"在数据库 {db_path} 中找到 {len(models)} 个 AI 模型记录。")
        
        # 检查并修复每个记录
        for model in models:
            model_id = model['id']
            model_name = model['name']
            
            # 检查是否有 stream_enabled 值
            try:
                stream_enabled = model['stream_enabled']
                if stream_enabled is None:
                    raise KeyError("stream_enabled is None")
                print(f"模型 '{model_name}' (ID: {model_id}) 已有 stream_enabled 值: {stream_enabled}")
            except (KeyError, IndexError):
                print(f"为模型 '{model_name}' (ID: {model_id}) 设置默认 stream_enabled 值")
                cursor.execute(
                    "UPDATE ai_models SET stream_enabled = 1 WHERE id = ?",
                    (model_id,)
                )
                conn.commit()
    
    # 关闭连接
    conn.close()
    return True

def fix_ai_models():
    """查找并修复所有可能的数据库中的 ai_models 表"""
    # 查找可能的数据库文件
    db_files = find_database_files()
    
    if not db_files:
        print("错误: 未找到任何数据库文件")
        sys.exit(1)
    
    print(f"找到 {len(db_files)} 个可能的数据库文件:")
    for i, db_file in enumerate(db_files):
        print(f"{i+1}. {db_file}")
    
    # 尝试修复每个数据库
    success = False
    for db_file in db_files:
        if fix_ai_models_in_db(db_file):
            success = True
    
    # 安装 requests 库（如果需要）
    try:
        import requests
        print("requests 库已安装。")
    except ImportError:
        print("requests 库未安装，尝试安装...")
        import subprocess
        try:
            subprocess.check_call([sys.executable, "-m", "pip", "install", "requests"])
            print("requests 库安装成功！")
        except Exception as e:
            print(f"安装 requests 库失败: {str(e)}")
            print("请手动运行: pip install requests")
    
    if success:
        print("修复完成！至少有一个数据库被成功修复。")
    else:
        print("警告: 未能修复任何数据库。")
        sys.exit(1)

if __name__ == "__main__":
    fix_ai_models() 